using System;
using System.IO;
using System.ServiceModel.Channels;
using System.Text;
using System.Xml;

namespace WcfExtensions.ServiceModel.Channels
{
    internal class MessageCompressor
    {
        private readonly CompressionAlgorithm _algorithm;
        private readonly CompressionLevel _level;

        public MessageCompressor(CompressionAlgorithm algorithm, CompressionLevel level)
        {
            _algorithm = algorithm;
            _level = level;
        }

        private CompressionAlgorithm Algorithm
        {
            get { return _algorithm; }
        }

        private CompressionLevel Level
        {
            get { return _level; }
        }

        public Message CompressMessage(Message toCompress)
        {
            if (Level == CompressionLevel.NoCompression)
            {
                return toCompress;
            }

            byte[] buffer = ExtractMessageBody(toCompress);
            if (buffer.Length == 0)
            {
                Message emptyMessage = CopyMessageAndReplaceContent(toCompress);
                toCompress.Close();
                return emptyMessage;
            }

            byte[] compressedBuf = DataCompressor.Compress(buffer, Algorithm, Level);

            var newMessageXml = new StringWriter();
            using (XmlWriter xmlWriter = XmlWriter.Create(newMessageXml))
            {
                WriteCompressedBody(xmlWriter, compressedBuf);
            }

            Message newMsg = CopyMessageAndReplaceContent(toCompress, newMessageXml.ToString());
            newMsg.Headers.Add(new CompressionMessageHeader(Algorithm, Level));
            toCompress.Close();

            return newMsg;
        }

        public static Message DecompressMessage(Message toDecompress, out CompressionMessageHeader compressionHeader)
        {
            compressionHeader = null;

            int index =
                toDecompress.Headers.FindHeader(WSCompressionStrings.Compression,
                                                WSCompressionStrings.NamespaceUri);
            if (index == -1)
                return toDecompress;

            using (XmlReader reader = toDecompress.Headers.GetReaderAtHeader(index))
            {
                compressionHeader = CompressionMessageHeader.ReadFrom(reader);
            }

            toDecompress.Headers.RemoveAll(WSCompressionStrings.Compression,
                                           WSCompressionStrings.NamespaceUri);

            string newMessageXml = DecompressMessageBody(toDecompress, compressionHeader);

            //Do not close the toDecompress message here, the framework manages its lifecycle.
            return CopyMessageAndReplaceContent(toDecompress, newMessageXml);
        }

        public static Message DecompressMessage(Message toDecompress)
        {
            CompressionMessageHeader header;
            return DecompressMessage(toDecompress, out header);
        }

        private static Message CopyMessageAndReplaceContent(Message sourceMessage, string newMessageXml)
        {
            XmlReader xmlReader = XmlReader.Create(new StringReader(newMessageXml));
            Message newMsg = Message.CreateMessage(sourceMessage.Version, null, xmlReader);
            newMsg.Headers.CopyHeadersFrom(sourceMessage);
            newMsg.Properties.CopyProperties(sourceMessage.Properties);
            return newMsg;
        }

        private static Message CopyMessageAndReplaceContent(Message sourceMessage)
        {
            Message newMsg = Message.CreateMessage(sourceMessage.Version, null);
            newMsg.Headers.CopyHeadersFrom(sourceMessage);
            newMsg.Properties.CopyProperties(sourceMessage.Properties);
            return newMsg;
        }

        private static byte[] ExtractMessageBody(Message toCompress)
        {
            using (XmlDictionaryReader reader = toCompress.GetReaderAtBodyContents())
            {
                return Encoding.UTF8.GetBytes(reader.ReadOuterXml());
            }
        }

        private static string DecompressMessageBody(Message toDecompress,
                                                         CompressionMessageHeader compressionHeader)
        {
            byte[] compressedData = ReadCompressedBody(toDecompress);

            byte[] data =
                DataCompressor.Decompress(compressedData, compressionHeader.Algorithm, compressionHeader.Level);

            return Encoding.UTF8.GetString(data);
        }

        private static byte[] ReadCompressedBody(Message message)
        {
            using (XmlReader reader = message.GetReaderAtBodyContents())
            {
                return Convert.FromBase64String(
                    reader.ReadElementString(WSCompressionStrings.CompressedBody,
                                             WSCompressionStrings.NamespaceUri));
            }
        }

        private static void WriteCompressedBody(XmlWriter xmlWriter, byte[] compressedData)
        {
            xmlWriter.WriteStartElement(WSCompressionStrings.CompressedBody,
                                        WSCompressionStrings.NamespaceUri);
            xmlWriter.WriteBase64(compressedData, 0, compressedData.Length);
            xmlWriter.WriteEndElement();
        }
    }
}